import os
import h5py
import numpy as np

def h5_to_npz(h5_path, npz_path=None):
    """把h5文件转为npz格式"""
    with h5py.File(h5_path, "r") as f:
        states = f["state"][:]
        obs = f["obs"][:]
        avails = f["avail_actions"][:]
        actions = f["actions"][:]
        rewards = f["reward"][:]
        dones = f["terminated"][:]
        actives = f["filled"][:]

    # ---- 关键修改：压缩掉最后一维 (N, T, 1) -> (N, T) ----
    if rewards.ndim == 3 and rewards.shape[-1] == 1:
        rewards = rewards.squeeze(-1)   # 去掉最后一维
    
    if dones.ndim == 3 and dones.shape[-1] == 1:
        dones = dones.squeeze(-1)

    if actives.ndim == 3 and actives.shape[-1] == 1:
        actives = actives.squeeze(-1)

    data = {
        "states": states,
        "obs": obs,
        "avails": avails,
        "actions": actions,
        "rewards": rewards,
        "dones": dones,
        "actives": actives,
    }

    if npz_path is None:
        npz_path = h5_path.replace(".h5", ".npz")
        # npz_path = npz_path.replace("offline_data", "offline_data")  # 保持目录一致

    os.makedirs(os.path.dirname(npz_path), exist_ok=True)
    np.savez_compressed(npz_path, **data)
    print(f"✅ 转换完成: {h5_path} -> {npz_path}")


if __name__ == "__main__":
    data_dir = "offline_data"
    out_dir = "offline_data"

    os.makedirs(out_dir, exist_ok=True)

    for fname in os.listdir(data_dir):
        if fname.endswith(".h5"):
            h5_path = os.path.join(data_dir, fname)
            npz_path = os.path.join(out_dir, fname.replace(".h5", ".npz"))
            h5_to_npz(h5_path, npz_path)
